{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state= 150 150\n",
      "action= [5, 3]\n",
      "reward= [0, 0]\n",
      "next_state= 150 150\n",
      "over= [False, False]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/pt39/lib/python3.9/site-packages/gym/spaces/box.py:127: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  logger.warn(f\"Box bound precision lowered by casting to {self.dtype}\")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<combat.Combat at 0x7fea37db4d00>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from combat import Combat\n",
    "\n",
    "\n",
    "def test_env():\n",
    "    state = env.reset()\n",
    "    action = env.action_space.sample()\n",
    "    next_state, reward, over, _ = env.step(action)\n",
    "\n",
    "    print('state=', len(state[0]), len(state[1]))\n",
    "    print('action=', action)\n",
    "    print('reward=', reward)\n",
    "    print('next_state=', len(next_state[0]), len(next_state[1]))\n",
    "    print('over=', over)\n",
    "\n",
    "\n",
    "env = Combat(grid_shape=(15, 15), n_agents=2, n_opponents=2)\n",
    "\n",
    "test_env()\n",
    "\n",
    "env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "import torch\n",
    "\n",
    "\n",
    "class PPO:\n",
    "    def __init__(self):\n",
    "        self.model_action = torch.nn.Sequential(\n",
    "            torch.nn.Linear(150, 64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(64, 64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(64, 7),\n",
    "            torch.nn.Softmax(dim=1),\n",
    "        )\n",
    "        self.model_value = torch.nn.Sequential(\n",
    "            torch.nn.Linear(150, 64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(64, 64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(64, 1),\n",
    "        )\n",
    "        self.optimizer_action = torch.optim.Adam(\n",
    "            self.model_action.parameters(), lr=3e-4)\n",
    "        self.optimizer_value = torch.optim.Adam(self.model_value.parameters(),\n",
    "                                                lr=3e-3)\n",
    "\n",
    "        self.mse_loss = torch.nn.MSELoss()\n",
    "\n",
    "    def get_action(self, state):\n",
    "        state = torch.FloatTensor(state).reshape(1, 150)\n",
    "\n",
    "        #[1, 150] -> [7]\n",
    "        weights = self.model_action(state).squeeze(dim=0).tolist()\n",
    "\n",
    "        #[7] -> scala\n",
    "        action = random.choices(range(7), weights=weights, k=1)[0]\n",
    "\n",
    "        return action\n",
    "\n",
    "    def _get_advantages(self, deltas):\n",
    "        advantages = []\n",
    "\n",
    "        #反向遍历deltas\n",
    "        s = 0.0\n",
    "        for delta in deltas[::-1]:\n",
    "            s = 0.99 * 0.97 * s + delta\n",
    "            advantages.append(s)\n",
    "\n",
    "        #逆序\n",
    "        advantages.reverse()\n",
    "        return advantages\n",
    "\n",
    "    def _get_target(self, next_state, reward, over):\n",
    "        #[b, 150] -> [b, 1]\n",
    "        target = self.model_value(next_state)\n",
    "        target *= 0.99\n",
    "        target *= (1 - over)\n",
    "        target += reward\n",
    "        return target\n",
    "\n",
    "    def _get_value(self, state):\n",
    "        #[b, 150] -> [b, 1]\n",
    "        return self.model_value(state)\n",
    "\n",
    "    def train(self, state, action, reward, next_state, over):\n",
    "        #state -> [b, 150]\n",
    "        #action -> [b, 1]\n",
    "        #reward -> [b, 1]\n",
    "        #next_state -> [b, 150]\n",
    "        #over -> [b, 1]\n",
    "\n",
    "        #[b, 1]\n",
    "        target = self._get_target(next_state, reward, over).detach()\n",
    "        #[b, 150] -> [b, 1]\n",
    "        value = self._get_value(state)\n",
    "\n",
    "        #[b, 1] - [b, 1] -> [b, 1] -> [b]\n",
    "        delta = (target - value).squeeze(dim=1).tolist()\n",
    "        #[b] -> [b]\n",
    "        advantages = self._get_advantages(delta)\n",
    "        #[b] -> [b, 1]\n",
    "        advantages = torch.FloatTensor(advantages).reshape(-1, 1)\n",
    "\n",
    "        #[b, 150] -> [b, 7]\n",
    "        old_prob = self.model_action(state)\n",
    "        #[b, 7] -> [b, 1]\n",
    "        old_prob = old_prob.gather(1, action)\n",
    "        #[b, 1] -> [b, 1]\n",
    "        old_prob = old_prob.log().detach()\n",
    "\n",
    "        for _ in range(1):\n",
    "            #[b, 150] -> [b, 7]\n",
    "            new_prob = self.model_action(state)\n",
    "            #[b, 7] -> [b, 1]\n",
    "            new_prob = new_prob.gather(1, action)\n",
    "            #[b, 1] -> [b, 1]\n",
    "            new_prob = new_prob.log()\n",
    "\n",
    "            #[b, 1] - [b, 1] -> [b, 1]\n",
    "            ratio = (new_prob - old_prob).exp()\n",
    "\n",
    "            #[b, 1] * [b, 1] -> [b, 1]\n",
    "            surr1 = ratio * advantages\n",
    "            #[b, 1] * [b, 1] -> [b, 1]\n",
    "            surr2 = torch.clamp(ratio, 0.8, 1.2) * advantages\n",
    "\n",
    "            #[b, 1]\n",
    "            loss_action = torch.min(surr1, surr2)\n",
    "            loss_action = -loss_action\n",
    "            #[b, 1] -> scala\n",
    "            loss_action = loss_action.mean()\n",
    "\n",
    "            self.optimizer_action.zero_grad()\n",
    "            loss_action.backward()\n",
    "            self.optimizer_action.step()\n",
    "\n",
    "            #[b, 4] -> [b, 1]\n",
    "            value = self._get_value(state)\n",
    "\n",
    "            self.optimizer_value.zero_grad()\n",
    "            #[b, 1],[b, 1] -> scala\n",
    "            loss_action = self.mse_loss(value, target)\n",
    "            self.optimizer_value.step()\n",
    "\n",
    "\n",
    "ppo = PPO()\n",
    "\n",
    "ppo.train(\n",
    "    torch.randn(5, 150),\n",
    "    torch.ones(5, 1).long(),\n",
    "    torch.randn(5, 1),\n",
    "    torch.randn(5, 150),\n",
    "    torch.zeros(5, 1).long(),\n",
    ")\n",
    "\n",
    "ppo.get_action(list(range(150)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'state': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]]),\n",
       "  'action': tensor([[2],\n",
       "          [3],\n",
       "          [6],\n",
       "          [5],\n",
       "          [5],\n",
       "          [0],\n",
       "          [2],\n",
       "          [1],\n",
       "          [1],\n",
       "          [0],\n",
       "          [2],\n",
       "          [5],\n",
       "          [3],\n",
       "          [4],\n",
       "          [4],\n",
       "          [2],\n",
       "          [0],\n",
       "          [2],\n",
       "          [5],\n",
       "          [1],\n",
       "          [1],\n",
       "          [5],\n",
       "          [5],\n",
       "          [4],\n",
       "          [1],\n",
       "          [5],\n",
       "          [2],\n",
       "          [0],\n",
       "          [2],\n",
       "          [5],\n",
       "          [1],\n",
       "          [2],\n",
       "          [1],\n",
       "          [6],\n",
       "          [4],\n",
       "          [4],\n",
       "          [0],\n",
       "          [1],\n",
       "          [0],\n",
       "          [2],\n",
       "          [2],\n",
       "          [2],\n",
       "          [0],\n",
       "          [1],\n",
       "          [6],\n",
       "          [3],\n",
       "          [6],\n",
       "          [6],\n",
       "          [0],\n",
       "          [2],\n",
       "          [6],\n",
       "          [0],\n",
       "          [2],\n",
       "          [6],\n",
       "          [5],\n",
       "          [2],\n",
       "          [2],\n",
       "          [4],\n",
       "          [6],\n",
       "          [0],\n",
       "          [3],\n",
       "          [3],\n",
       "          [0],\n",
       "          [1],\n",
       "          [0],\n",
       "          [3],\n",
       "          [2],\n",
       "          [4],\n",
       "          [2],\n",
       "          [3],\n",
       "          [4],\n",
       "          [2],\n",
       "          [0],\n",
       "          [2],\n",
       "          [1],\n",
       "          [0],\n",
       "          [5],\n",
       "          [2],\n",
       "          [4],\n",
       "          [5],\n",
       "          [1],\n",
       "          [5],\n",
       "          [5],\n",
       "          [1],\n",
       "          [0],\n",
       "          [5],\n",
       "          [2],\n",
       "          [4],\n",
       "          [1],\n",
       "          [6],\n",
       "          [5],\n",
       "          [1],\n",
       "          [2],\n",
       "          [1],\n",
       "          [0],\n",
       "          [2],\n",
       "          [2],\n",
       "          [5],\n",
       "          [5],\n",
       "          [6]]),\n",
       "  'reward': tensor([[-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000]]),\n",
       "  'next_state': tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "          [0., 0., 0.,  ..., 0., 0., 0.]]),\n",
       "  'over': tensor([[0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0]])},\n",
       " {'state': tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          ...,\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.1333],\n",
       "          [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000, -1.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]),\n",
       "  'action': tensor([[4],\n",
       "          [0],\n",
       "          [4],\n",
       "          [1],\n",
       "          [1],\n",
       "          [2],\n",
       "          [4],\n",
       "          [5],\n",
       "          [4],\n",
       "          [0],\n",
       "          [5],\n",
       "          [5],\n",
       "          [0],\n",
       "          [0],\n",
       "          [6],\n",
       "          [3],\n",
       "          [0],\n",
       "          [6],\n",
       "          [3],\n",
       "          [2],\n",
       "          [4],\n",
       "          [3],\n",
       "          [6],\n",
       "          [0],\n",
       "          [5],\n",
       "          [1],\n",
       "          [4],\n",
       "          [2],\n",
       "          [5],\n",
       "          [1],\n",
       "          [6],\n",
       "          [6],\n",
       "          [1],\n",
       "          [3],\n",
       "          [5],\n",
       "          [3],\n",
       "          [3],\n",
       "          [2],\n",
       "          [0],\n",
       "          [5],\n",
       "          [3],\n",
       "          [6],\n",
       "          [6],\n",
       "          [0],\n",
       "          [1],\n",
       "          [5],\n",
       "          [0],\n",
       "          [6],\n",
       "          [4],\n",
       "          [4],\n",
       "          [5],\n",
       "          [0],\n",
       "          [1],\n",
       "          [3],\n",
       "          [0],\n",
       "          [5],\n",
       "          [4],\n",
       "          [3],\n",
       "          [4],\n",
       "          [3],\n",
       "          [2],\n",
       "          [2],\n",
       "          [1],\n",
       "          [4],\n",
       "          [4],\n",
       "          [4],\n",
       "          [3],\n",
       "          [5],\n",
       "          [3],\n",
       "          [1],\n",
       "          [4],\n",
       "          [2],\n",
       "          [5],\n",
       "          [6],\n",
       "          [2],\n",
       "          [2],\n",
       "          [6],\n",
       "          [2],\n",
       "          [0],\n",
       "          [0],\n",
       "          [5],\n",
       "          [6],\n",
       "          [5],\n",
       "          [1],\n",
       "          [2],\n",
       "          [0],\n",
       "          [1],\n",
       "          [0],\n",
       "          [6],\n",
       "          [4],\n",
       "          [6],\n",
       "          [0],\n",
       "          [3],\n",
       "          [5],\n",
       "          [5],\n",
       "          [2],\n",
       "          [5],\n",
       "          [2],\n",
       "          [1],\n",
       "          [1]]),\n",
       "  'reward': tensor([[-0.1000],\n",
       "          [-0.1000],\n",
       "          [-1.1000],\n",
       "          [-1.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-1.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000],\n",
       "          [-0.1000]]),\n",
       "  'next_state': tensor([[ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
       "          [ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
       "          [ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
       "          ...,\n",
       "          [ 0.,  0.,  0.,  ...,  0.,  0.,  0.],\n",
       "          [ 0., -1.,  0.,  ...,  0.,  0.,  0.],\n",
       "          [ 0.,  0.,  0.,  ...,  0.,  0.,  0.]]),\n",
       "  'over': tensor([[0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0],\n",
       "          [0]])},\n",
       " False)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_data():\n",
    "    data0 = {\n",
    "        'state': [],\n",
    "        'action': [],\n",
    "        'reward': [],\n",
    "        'next_state': [],\n",
    "        'over': [],\n",
    "    }\n",
    "\n",
    "    data1 = {\n",
    "        'state': [],\n",
    "        'action': [],\n",
    "        'reward': [],\n",
    "        'next_state': [],\n",
    "        'over': [],\n",
    "    }\n",
    "\n",
    "    state = env.reset()\n",
    "    over = False\n",
    "    while not over:\n",
    "        action = [None, None]\n",
    "\n",
    "        action[0] = ppo.get_action(state[0])\n",
    "        action[1] = ppo.get_action(state[1])\n",
    "\n",
    "        next_state, reward, over, info = env.step(action)\n",
    "        win = info['win']\n",
    "        del info\n",
    "\n",
    "        #对reward进行偏移\n",
    "        if win:\n",
    "            reward[0] += 100\n",
    "            reward[1] += 100\n",
    "        else:\n",
    "            reward[0] -= 0.1\n",
    "            reward[1] -= 0.1\n",
    "\n",
    "        data0['state'].append(state[0])\n",
    "        data0['action'].append(action[0])\n",
    "        data0['reward'].append(reward[0])\n",
    "        data0['next_state'].append(next_state[0])\n",
    "        data0['over'].append(False)  #常量\n",
    "\n",
    "        data1['state'].append(state[1])\n",
    "        data1['action'].append(action[1])\n",
    "        data1['reward'].append(reward[1])\n",
    "        data1['next_state'].append(next_state[1])\n",
    "        data1['over'].append(False)  #常量\n",
    "\n",
    "        state = next_state\n",
    "        over = over[0] and over[1]\n",
    "\n",
    "    data0['state'] = torch.FloatTensor(data0['state']).reshape(-1, 150)\n",
    "    data0['action'] = torch.LongTensor(data0['action']).reshape(-1, 1)\n",
    "    data0['reward'] = torch.FloatTensor(data0['reward']).reshape(-1, 1)\n",
    "    data0['next_state'] = torch.FloatTensor(data0['next_state']).reshape(-1, 150)\n",
    "    data0['over'] = torch.LongTensor(data0['over']).reshape(-1, 1)\n",
    "\n",
    "    data1['state'] = torch.FloatTensor(data1['state']).reshape(-1, 150)\n",
    "    data1['action'] = torch.LongTensor(data1['action']).reshape(-1, 1)\n",
    "    data1['reward'] = torch.FloatTensor(data1['reward']).reshape(-1, 1)\n",
    "    data1['next_state'] = torch.FloatTensor(data1['next_state']).reshape(-1, 150)\n",
    "    data1['over'] = torch.LongTensor(data1['over']).reshape(-1, 1)\n",
    "\n",
    "    return data0, data1, win\n",
    "\n",
    "\n",
    "get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 10107,
     "status": "ok",
     "timestamp": 1650012696153,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "-_L_dhppItIk",
    "outputId": "6c1eecf0-fd72-4d13-ad05-192463636129"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0.0\n",
      "10000 0.1\n",
      "20000 0.21\n",
      "30000 0.27\n",
      "40000 0.29\n",
      "50000 0.28\n",
      "60000 0.29\n",
      "70000 0.13\n",
      "80000 0.27\n",
      "90000 0.23\n",
      "100000 0.23\n",
      "110000 0.26\n",
      "120000 0.37\n",
      "130000 0.34\n",
      "140000 0.29\n",
      "150000 0.27\n",
      "160000 0.33\n",
      "170000 0.28\n",
      "180000 0.3\n",
      "190000 0.37\n"
     ]
    }
   ],
   "source": [
    "wins = []\n",
    "for i in range(200000):\n",
    "    data0, data1, win = get_data()\n",
    "    wins.append(win)\n",
    "\n",
    "    ppo.train(**data0)\n",
    "    ppo.train(**data1)\n",
    "\n",
    "    if i % 10000 == 0:\n",
    "        wins = wins[-100:]\n",
    "        print(i, sum(wins) / len(wins))\n",
    "        wins = []"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "第20章-多智能体强化学习入门.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
